/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2022-2022. All rights reserved.
 * Description: 集合通信算子build graph
 * Author: lilianlin
 * Create: 2022-07-15
 */

#ifndef HCOM_BUILD_GRAPH_H
#define HCOM_BUILD_GRAPH_H
#include "exe_graph/lowering/dev_mem_value_holder.h"
#include "graph/node.h"
#include "comm.h"
#include "register/node_converter_registry.h"

namespace hccl {

constexpr u32 LAUNCH_RECV_KERNEL_OUTPUT_SIZE = 2; /* recv kerner的输出 */

enum class HcomOpType {
    HCOM_ALL_GATHER,
    HCOM_ALL_GATHER_V,
    HCOM_ALL_REDUCE,
    HCOM_BROADCAST,
    HCOM_REDUCE_SCATTER,
    HCOM_REDUCE_SCATTER_V,
    HCOM_ALL_TO_ALL_V,
    HCOM_ALL_TO_ALL_VC,
    HCOM_GATHER_ALL_TO_ALL_V,
    HCOM_REMOTE_READ,
    HCOM_REMOTE_REF_READ,
    HCOM_REMOTE_WRITE,
    HCOM_REMOTE_SCATTER_WRITE,
    HCOM_SEND,
    HCOM_REDUCE,
    COLL_REMOTE_LOOK_UP_PAIRED,
    COLL_REMOTE_UPDATE_PAIRED,
    COLL_REMOTE_LOOK_UP_UNIQUED_PAIRED,
    REMOTE_LOOK_UP,
    HCOM_RECEIVE,
};

struct HcomAllReduceOpAttr {
    HcclReduceOp reduction;
};

struct HcomAllGatherOpAttr {
    int32_t rankSize;
};

struct HcomAllGatherVOpAttr {
    HcclDataType recvType;
};

struct HcomBroadcastOpAttr {
    int32_t root;
};

struct HcomReduceScatterOpAttr {
    HcclReduceOp reduction;
    int32_t rankSize;
};

struct HcomReduceScatterVOpAttr {
    HcclReduceOp reduction;
    int32_t rankSize;
};

struct HcomAllToAllVOpAttr {
    HcclDataType recvType;
};

struct HcomGatherAllToAllVOpAttr {
    int32_t addrLength;
};

struct HcomRecvOpAttr {
    int32_t srcRank;
    int32_t srTag;
};

struct HcomSendOpAttr {
    int32_t destRank;
    int32_t srTag;
};

struct HcomReduceOpAttr {
    int32_t root;
    HcclReduceOp reduction;
};

struct HcomRemoteLookupOpAttr {
    int32_t tag;
    int32_t insert_option;
    int32_t max_num;
    int32_t embedding_dim;
};

struct HcomCollRemoteLookupPairedOpAttr {
    int32_t tag;
    int32_t insert_option;
    int32_t max_num;
    int32_t embedding_dim;
    int32_t flags;
};

struct HcomCollRemoteUpdatePairedOpAttr {
    int32_t tag;
    int32_t max_num;
    int32_t embedding_dim;
};

struct HcomCollRemoteLookupPairedAndUniqueOpAttr {
    int32_t tag;
    int32_t insert_option;
    int32_t max_num;
    int32_t embedding_dim;
    int32_t flags;
};
struct HcomOpAttr {
    HcomOpType opType;
    HcclDataType dataType;
    u32 qosConfig;
    char group[GROUP_NAME_MAX_LEN + 1];
    union {
        HcomAllGatherOpAttr allgather;
        HcomAllGatherVOpAttr allgatherv;
        HcomAllReduceOpAttr allreduce;
        HcomBroadcastOpAttr broadcast;
        HcomReduceScatterOpAttr reducescatter;
        HcomReduceScatterVOpAttr reducescatterv;
        HcomAllToAllVOpAttr alltoallv;
        HcomGatherAllToAllVOpAttr gatheralltoallv;
        HcomSendOpAttr send;
        HcomRecvOpAttr recv;
        HcomReduceOpAttr reduce;
        HcomRemoteLookupOpAttr remotelookup;
        HcomCollRemoteLookupPairedOpAttr collremotelookuppaired;
        HcomCollRemoteUpdatePairedOpAttr collremoteupdatepaired;
        HcomCollRemoteLookupPairedAndUniqueOpAttr collremotelookuppairedandunique;
    } op;
};


struct HcomLaunchArg {
    gert::bg::ValueHolderPtr opArgs;
    gert::bg::ValueHolderPtr stream;
};

gert::bg::ValueHolderPtr GenerateHcomOpArgs(const ge::NodePtr &node);

std::vector<gert::bg::ValueHolderPtr> LaunchHcomOpKernel(const HcomLaunchArg &opAttr,
    const std::vector<gert::bg::DevMemValueHolderPtr> &inputAddrs,
    const std::vector<gert::bg::DevMemValueHolderPtr> &outputAddrs,
    const std::vector<gert::bg::ValueHolderPtr> &inputShapes,
    const std::vector<gert::bg::ValueHolderPtr> &outputShapes);

std::vector<gert::bg::DevMemValueHolderPtr> LaunchRecvOpKernel(const HcomLaunchArg &args, const ge::NodePtr &node,
    const gert::LowerInput &lowerInput);
}  // namespace hccl
#endif // HCOM_BUILD_GRAPH_H